{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "from pytorch3d.structures import Pointclouds\n",
    "from pytorch3d.vis.plotly_vis import plot_scene\n",
    "from pytorch3d.renderer.cameras import PerspectiveCameras\n",
    "\n",
    "import torch\n",
    "import numpy as np\n",
    "from read_write_model import read_model, qvec2rotmat"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# select sequence tag to show here\n",
    "\n",
    "# tag = \"apple_189_20393_38136\"\n",
    "tag = \"teddybear_34_1479_4753\"\n",
    "# tag = \"toytrain_240_25394_51994\"\n",
    "\n",
    "path = f\"../var/colmap/{tag}/model/sparse/0\"\n",
    "\n",
    "# read COLMAP reconstruction model\n",
    "cameras, images, points3D = read_model(path, ext = \".bin\")\n",
    "\n",
    "print(f\"{len(cameras)} cameras found\")\n",
    "print(f\"{len(images)} poses found\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# extract point cloud\n",
    "rgb = np.array([v.rgb for v in points3D.values()])\n",
    "xyz = np.array([v.xyz for v in points3D.values()])\n",
    "\n",
    "rgb = torch.tensor(rgb, dtype=torch.uint8)\n",
    "xyz = torch.tensor(xyz, dtype=torch.float32)\n",
    "\n",
    "print(f\"#{rgb.shape[0]} points found\")\n",
    "\n",
    "point_cloud = Pointclouds(points=[xyz], features=[rgb/255.])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# extract camera poses\n",
    "focal_lengths = []\n",
    "principal_points = []\n",
    "R = []\n",
    "T = []\n",
    "\n",
    "for k, image in images.items():\n",
    "    # intrinsics don't matter much for visualization\n",
    "    focal_lengths.append((1., 1.))\n",
    "    principal_points.append((0., 0.))\n",
    "\n",
    "    # extract rotation and translation\n",
    "    R.append(qvec2rotmat(image.qvec))\n",
    "    T.append(image.tvec)\n",
    "\n",
    "cams = PerspectiveCameras(\n",
    "    focal_length=focal_lengths,\n",
    "    principal_point=principal_points,\n",
    "    R=R,\n",
    "    T=T,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "plot_scene(\n",
    "    {\"COLMAP Reconstruction (using SiLK keypoints)\": {tag: point_cloud[0], \"cameras\": cams}},\n",
    "    pointcloud_max_points=rgb.shape[0],\n",
    "    raybundle_max_rays=rgb.shape[0],\n",
    "    raybundle_max_points_per_ray=rgb.shape[0],\n",
    "    width=400,\n",
    "    height=400,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
